{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06e37c53-a3e6-4753-987d-e20d1f5278e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import Tool, AgentExecutor, LLMSingleActionAgent, AgentOutputParser\n",
    "from langchain.prompts import StringPromptTemplate\n",
    "from langchain.llms import OpenAI\n",
    "from langchain.utilities import SerpAPIWrapper\n",
    "from langchain.chains import LLMChain\n",
    "from typing import List, Union\n",
    "from langchain.schema import AgentAction, AgentFinish, OutputParserException\n",
    "import re\n",
    "# Define which tools the agent can use to answer user queries\n",
    "import os\n",
    "os.environ[\"SERPAPI_API_KEY\"] = 'd329aa9848c97a69d8ad344dddf28af5dd8ab50e7ea1aa6271cb8c1ee94f4417'\n",
    " \n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b862f201-1ee8-4da4-919f-f35946b20462",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the base template\n",
    "template = \"\"\"Answer the following questions as best you can, but speaking as a pirate might speak. You have access to the following tools:\n",
    "\n",
    "{tools}\n",
    "\n",
    "Use the following format:\n",
    "\n",
    "Question: the input question you must answer\n",
    "Thought: you should always think about what to do\n",
    "Action: the action to take, should be one of [{tool_names}]\n",
    "Action Input: the input to the action\n",
    "Observation: the result of the action\n",
    "... (this Thought/Action/Action Input/Observation can repeat N times)\n",
    "Thought: I now know the final answer\n",
    "Final Answer: the final answer to the original input question\n",
    "\n",
    "Begin! Remember to speak as a pirate when giving your final answer. Use lots of \"Arg\"s\n",
    "\n",
    "\n",
    "New question: {input}\n",
    "{agent_scratchpad}\"\"\"\n",
    "\n",
    "# Set up a prompt template\n",
    "class CustomPromptTemplate(StringPromptTemplate):\n",
    "    # The template to use\n",
    "    template: str\n",
    "    # The list of tools available\n",
    "    tools: List[Tool]\n",
    "\n",
    "    def format(self, **kwargs) -> str:\n",
    "        # Get the intermediate steps (AgentAction, Observation tuples)\n",
    "        # Format them in a particular way\n",
    "        intermediate_steps = kwargs.pop(\"intermediate_steps\")\n",
    "        thoughts = \"\"\n",
    "        for action, observation in intermediate_steps:\n",
    "            thoughts += action.log\n",
    "            thoughts += f\"\\nObservation: {observation}\\nThought: \"\n",
    "        # Set the agent_scratchpad variable to that value\n",
    "        kwargs[\"agent_scratchpad\"] = thoughts\n",
    "        # Create a tools variable from the list of tools provided\n",
    "        kwargs[\"tools\"] = \"\\n\".join([f\"{tool.name}: {tool.description}\" for tool in self.tools])\n",
    "        # Create a list of tool names for the tools provided\n",
    "        kwargs[\"tool_names\"] = \", \".join([tool.name for tool in self.tools])\n",
    "        return self.template.format(**kwargs)\n",
    "class CustomOutputParser(AgentOutputParser):\n",
    "\n",
    "    def parse(self, llm_output: str) -> Union[AgentAction, AgentFinish]:\n",
    "        # Check if agent should finish\n",
    "        if \"Final Answer:\" in llm_output:\n",
    "            return AgentFinish(\n",
    "                # Return values is generally always a dictionary with a single `output` key\n",
    "                # It is not recommended to try anything else at the moment :)\n",
    "                return_values={\"output\": llm_output.split(\"Final Answer:\")[-1].strip()},\n",
    "                log=llm_output,\n",
    "            )\n",
    "        # Parse out the action and action input\n",
    "        regex = r\"Action\\s*\\d*\\s*:(.*?)\\nAction\\s*\\d*\\s*Input\\s*\\d*\\s*:[\\s]*(.*)\"\n",
    "        match = re.search(regex, llm_output, re.DOTALL)\n",
    "        if not match:\n",
    "            raise OutputParserException(f\"Could not parse LLM output: `{llm_output}`\")\n",
    "        action = match.group(1).strip()\n",
    "        action_input = match.group(2)\n",
    "        # Return the action and action input\n",
    "        return AgentAction(tool=action, tool_input=action_input.strip(\" \").strip('\"'), log=llm_output)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9cd7eda-bfcb-445e-bdc6-d07e556aac20",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_parser = CustomOutputParser()\n",
    "from langchain.chains.llm_math.base import LLMMathChain\n",
    "openai_api_base = \"http://127.0.0.1:6003/v1\"\n",
    "openai_api_key = \"test\"\n",
    "llm = OpenAI(openai_api_base=openai_api_base, openai_api_key=openai_api_key,temperature=0, )\n",
    "tools = [\n",
    "    Tool(\n",
    "        name=\"Search\",\n",
    "        func=search.run,\n",
    "        description=\"useful for when you need to answer questions about current events\"\n",
    "    ),\n",
    "    Tool(\n",
    "        name=\"Calculator\",\n",
    "        description=\"Useful for when you need to answer questions about math.\",\n",
    "        func=LLMMathChain.from_llm(llm=llm).run,\n",
    "        \n",
    "    )\n",
    "]\n",
    "prompt = CustomPromptTemplate(\n",
    "    template=template,\n",
    "    tools=tools,\n",
    "    # This omits the `agent_scratchpad`, `tools`, and `tool_names` variables because those are generated dynamically\n",
    "    # This includes the `intermediate_steps` variable because that is needed\n",
    "    input_variables=[\"input\", \"intermediate_steps\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5ab5ac3-eeb8-4331-a1a2-fd65387d7420",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm_chain = LLMChain(llm=llm, prompt=prompt)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
